package com.atguigu.java.ai.langchain4j.store;


import com.atguigu.java.ai.langchain4j.bean.ChatMessages;
import com.atguigu.java.ai.langchain4j.bean.ChatMessages1;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Component;

import java.util.LinkedList;
import java.util.List;

/**
 * @Author: LZY
 * @Date: 2025/09/01 20:30
 * @Description:
 **/
@Component  // 标记这个类是一个Spring组件
public class MongoChatMemoryStore implements ChatMemoryStore {

    @Autowired
    //引入依赖  MongoTemplate 提供了海量的方法，几乎覆盖了所有 MongoDB 操作
    private MongoTemplate mongoTemplate;

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        Query query = new Query();
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        query.addCriteria(criteria);

        ChatMessages1 chatMessages = mongoTemplate.findOne(query, ChatMessages1.class);
        if (chatMessages == null) return new LinkedList<>();
        //
        List<ChatMessage> messagesFromJson = ChatMessageDeserializer.messagesFromJson(chatMessages.getContent());
        return messagesFromJson;
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        Query query = new Query();
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        query.addCriteria(criteria);

        Update update = new Update();
        update.set("content", ChatMessageSerializer.messagesToJson(messages));

        //根据query条件能查询出文档，则修改文档；否则新增文档
        mongoTemplate.upsert(query, update, ChatMessages.class);
    }

    @Override
    public void deleteMessages(Object memoryId) {
        Query query = new Query();
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        query.addCriteria(criteria);

        mongoTemplate.remove(query, ChatMessages.class);
    }
}
